import socket
import threading
import time
import signal  # 新增：用于处理信号

class TelnetTestServer:
    def __init__(self, host='0.0.0.0', port=23):
        self.host = host
        self.port = port
        self.server_socket = None
        self.running = False
        self.client_threads = []  # 新增：跟踪客户端线程
        self.lock = threading.Lock()  # 新增：线程安全锁

    def start(self):
        # 新增：注册信号处理函数，捕获Ctrl+C
        signal.signal(signal.SIGINT, self.handle_signal)
        signal.signal(signal.SIGTERM, self.handle_signal)

        self.running = True
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        # 新增：设置非阻塞模式，避免accept()无限阻塞
        self.server_socket.setblocking(False)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(5)
        print(f"Telnet test server started on {self.host}:{self.port}")
        print("Press Ctrl+C to stop server")
        
        try:
            while self.running:
                try:
                    # 非阻塞模式下，若没有新连接会抛BlockingIOError
                    client_socket, client_addr = self.server_socket.accept()
                    print(f"New connection from {client_addr}")
                    client_thread = threading.Thread(
                        target=self.handle_client,
                        args=(client_socket, client_addr),
                        daemon=True  # 新增：设置为守护线程，主程序退出时自动终止
                    )
                    with self.lock:
                        self.client_threads.append(client_thread)
                    client_thread.start()
                except BlockingIOError:
                    # 没有新连接时，短暂休眠避免CPU空转
                    time.sleep(0.1)
        finally:
            self.stop()

    def handle_client(self, client_socket, client_addr):
        try:
            # 发送欢迎消息
            client_socket.sendall(b"Welcome to the test Telnet server!\r\n")
            client_socket.sendall(b"Type 'exit' to disconnect\r\n")
            client_socket.sendall(b"Type 'info' to get server info\r\n")
            client_socket.sendall(b"> ")
            
            while self.running:  # 新增：检查服务器是否在运行
                # 客户端socket设置超时，避免无限阻塞
                client_socket.settimeout(1.0)
                try:
                    data = client_socket.recv(1024)
                except socket.timeout:
                    continue  # 超时后继续循环，检查服务器是否仍在运行
                
                if not data:
                    break
                
                data_str = data.decode('utf-8', errors='ignore').strip()
                print(f"Received from {client_addr}: {data_str}")
                
                if data_str.lower() == 'exit':
                    client_socket.sendall(b"Goodbye!\r\n")
                    break
                elif data_str.lower() == 'info':
                    response = f"Server info: Python Telnet Test Server running on {self.host}:{self.port}\r\n"
                    client_socket.sendall(response.encode('utf-8'))
                elif data_str.lower() == 'help':
                    client_socket.sendall(b"Commands: exit, info, help\r\n")
                else:
                    response = f"You sent: {data_str}\r\n"
                    client_socket.sendall(response.encode('utf-8'))
                
                client_socket.sendall(b"> ")
        except Exception as e:
            if self.running:  # 仅在服务器正常运行时打印错误
                print(f"Error handling client {client_addr}: {e}")
        finally:
            client_socket.close()
            print(f"Connection with {client_addr} closed")
            # 从线程列表中移除当前线程
            with self.lock:
                self.client_threads = [t for t in self.client_threads if t.is_alive()]

    def stop(self):
        self.running = False
        if self.server_socket:
            self.server_socket.close()
            print("Server socket closed")
        # 等待所有客户端线程终止（最多等待2秒）
        start_time = time.time()
        with self.lock:
            for thread in self.client_threads:
                while thread.is_alive() and (time.time() - start_time) < 2:
                    time.sleep(0.1)
        print("Server stopped successfully")

    def handle_signal(self, signum, frame):
        """信号处理函数：捕获Ctrl+C后触发服务器停止"""
        print("\nReceived stop signal (Ctrl+C), stopping server...")
        self.running = False

if __name__ == "__main__":
    server = TelnetTestServer(port=2323)
    server.start()